import os
os.environ['CUDA_VISIBLE_DEVICES']='0'

import torch,pdb
from diffusers import FluxPipeline, FluxPriorReduxPipeline
from diffusers.utils import load_image
from util_flux import process_img_1024,vertical_concat_images,horizontal_concat_images
from util_flux import resize_with_aspect
from PIL import Image
from itertools import product
from util_for_os import osj,ose

from MODEL_CKP import FLUX

pipe = FluxPipeline.from_pretrained(
    FLUX , 
    # text_encoder=None,
    # text_encoder_2=None,
    torch_dtype=torch.bfloat16
).to("cuda")

import random
def generate_image(text_prompt,height=1024,width=1024,
                   steps=20,scale=4.5,seed = -1):
    if seed == -1: seed = random.randint(0,10**10)
    print('seed:',seed)
    with torch.no_grad():
        image = pipe(
            prompt=text_prompt,
            height=height,
            width=width,
            num_inference_steps=steps,
            guidance_scale=scale,
            generator=torch.Generator().manual_seed(seed),
        ).images[0]

    image.save( f'tmp_dev/{text_prompt[:8]}_{seed}.jpg' )

    return image,seed

import gradio as gr
# 构建gradio，输入 提示词等内容 生成 图片
def start_gradio():
    # 创建Gradio界面
    with gr.Blocks(title="AI 图像生成器") as demo:
        gr.Markdown("# 🎨 AI 图像生成器")
        gr.Markdown("输入提示词，生成高质量图像")
        
        with gr.Row():
            with gr.Column():
                prompt = gr.Textbox(
                    label="提示词",
                    placeholder="请输入描述图像的文本...",
                    lines=3
                )
                with gr.Accordion("高级设置", open=False):
                    width = gr.Slider(
                        label="宽度",
                        minimum=512,
                        maximum=2048,
                        step=64,
                        value=1024
                    )
                    height = gr.Slider(
                        label="高度",
                        minimum=512,
                        maximum=2048,
                        step=64,
                        value=1024
                    )
                    steps = gr.Slider(
                        label="推理步数",
                        minimum=10,
                        maximum=50,
                        step=1,
                        value=20
                    )
                    scale = gr.Slider(
                        label="引导尺度",
                        minimum=1.0,
                        maximum=10.0,
                        step=0.5,
                        value=4.5
                    )
                    seed = gr.Number(
                        label="随机种子 (-1表示随机)",
                        value=-1
                    )
                submit_btn = gr.Button("生成图像", variant="primary")
            
            with gr.Column():
                output_image = gr.Image(label="生成的图像",height=512)
                seed_display = gr.Textbox(label="使用的种子")
        
        # 示例提示词
        examples = gr.Examples(
            examples=[
                ["一只穿着宇航服的柴犬在月球上漫步，科幻风格，4K高清"],
                ["未来城市夜景，赛博朋克风格，霓虹灯光，雨夜街道"],
                ["宁静的山水风景，中国水墨画风格，云雾缭绕"]
            ],
            inputs=[prompt, height, width, steps, scale, seed],
            label="示例提示词 (点击试试)"
        )
        
        # 绑定函数
        submit_btn.click(
            fn=generate_image,
            inputs=[prompt, height, width, steps, scale, seed],
            outputs=[output_image, seed_display]
        )

    demo.launch( server_name='0.0.0.0',server_port=20024 )

if __name__=='__main__':
    start_gradio()